import os
import sqlite3
from llama_index.core import Document
from llama_index.core import VectorStoreIndex
from pprint import pprint


def create_sample_database():
    # 连接到sample.db的SQLLite数据库，如果此文件不存在则会自动创建
    conn = sqlite3.connect("sample.db")
    # 获取数据库的游标对象
    cursor = conn.cursor()

    # 创建users表（如果不存在则新建）
    cursor.execute(
        """
        CREATE TABLE IF NOT EXISTS users (
            id INTEGER PRIMARY KEY,
            name TEXT NOT NULL,
            email TEXT,
            department TEXT,
            salary REAL
        )
        """
    )

    # 定义要插入的示例数据
    sample_data = [
        (1, "张三", "zhangsan@company.com", "技术部", 8000),
        (2, "李四", "lisi@company.com", "市场部", 7500),
        (3, "王五", "wangwu@company.com", "人事部", 7000),
        (4, "赵六", "zhaoliu@company.com", "技术部", 8500),
        (5, "钱七", "qianqi@company.com", "财务部", 7200),
    ]

    # 批量插入示例数据到users表
    cursor.executemany(
        """
        INSERT OR REPLACE INTO users (id, name, email, department, salary)
        VALUES (?, ?, ?, ?, ?)
    """,
        sample_data,
    )
    conn.commit()
    conn.close()
    print("数据库创建完成")


def load_data_from_database():
    # 连接到sample.db的SQLLite数据库，如果此文件不存在则会自动创建
    conn = sqlite3.connect("sample.db")
    # 获取数据库的游标对象
    cursor = conn.cursor()
    cursor.execute("SELECT * FROM users")
    # 获取所有的查询结果行
    rows = cursor.fetchall()
    documents = []
    for row in rows:
        doc = Document(
            text=f"id={row[0]},name={row[1]},email={row[2]},department={row[3]},salary={row[4]}",
            metadata={
                "id": row[0],
                "name": row[1],
                "email": row[2],
                "department": row[3],
                "salary": row[4],
                "source": "database",
            },
        )
        documents.append(doc)
    conn.close()
    return documents


# 创建示例数据库
create_sample_database()
# 从数据库里加载文档
documents = load_data_from_database()
index = VectorStoreIndex.from_documents(documents)
query_engine = index.as_query_engine()
# 查询1：技术部有哪些员工
print("查询1: 技术部有哪些员工？")
response1 = query_engine.query("技术部有哪些员工？")
print(f"回答: {response1}")
print()
# 查询2：薪资最高的员工是谁
print("查询2: 薪资最高的员工是谁？")
response2 = query_engine.query("薪资最高的员工是谁？")
print(f"回答: {response2}")
print()
# 查询3：市场部员工信息
print("查询3: 市场部的员工信息")
response3 = query_engine.query("市场部的员工信息")
print(f"回答: {response3}")
print()
# 查询4：公司的平均薪资是多少
print("查询4: 公司的平均薪资是多少？")
response4 = query_engine.query("公司的平均薪资是多少？")
print(f"回答: {response4}")
print()
